"""
	This file contains the DecisionTree class, handling the construction and
	traversal of a decision tree.

TODO:
Convert from old Dataset to new Dataset
Rework AttributeSelectors
Inject Logging
Unit Test
"""

import random
import AttributeSelectors

class DecisionTree:
	AttributeSelector = AttributeSelectors.iterative_dichotomiser_3

	def __init__(self, dataset):
		# Save off the dataset and attribute names
		self.Dataset = dataset
		self.TargetedAttributeValue = None

		# Select an attribute to divide on
		result = self.AttributeSelector(dataset, attributeNames)
		
		# Handling Terminal Nodes
		if len(dataset) == 0:
			self.Children = {}
			return
		if result == None:
			self.TargetedAttribute, self.InformationGain = None, None
			self.Children = {}
			return

		self.TargetedAttribute, self.InformationGain = result
		
		# Divide the dataset by the classifications of the selected attribute
		dividedDataset = {}
		i = attributeNames.index(self.TargetedAttribute)
		for item, result in dataset:
			if not (item[i] in dividedDataset):
				dividedDataset[item[i]] = []
			dividedDataset[item[i]].append((item[:i] + item[i+1:], result))

		# Remove the selected attribute from attribute names of descendent nodes
		newAttributes = attributeNames[:i] + attributeNames[i+1:]

		# If there aren't any more attributes to select on, become a terminal node
		if len(newAttributes) == 0:
			self.TargetedAttribute, self.InformationGain = None, None
			self.Children = {}
			return

		# Construct a Binary Decision Tree on each dataset
		children = {}
		for key in dividedDataset:
			children[key] = DecisionTree(newAttributes, dividedDataset[key])
		
		# Save off the children
		self.Children = children

		# Mark the children with the value of their selected attribute
		for key in self.Children:
			self.Children[key].TargetedAttributeValue = key
	
	def to_string(self):
		result = ""
		for item in self._to_string():
			result += item + "\n"
		return result

	def _to_string(self):
		if (self.TargetedAttribute, self.InformationGain) == (None, None):
			frequencies = {}
			for item in self.Dataset:
				if not (item[1] in frequencies):
					frequencies[item[1]] = 0.0
				frequencies[item[1]] += 1.0
			for item in frequencies:
				frequencies[item] = frequencies[item] / len(self.Dataset)
			result = ["Probability Distribution:{0}".format(frequencies)]
			return result
		
		if self.TargetedAttributeValue == None:
			result = [
			"{0} Gain: {1}".format(self.TargetedAttribute, self.InformationGain)
			]
		else:
			result = [
			"Parent Attribute Value:{0} SelectedAttribute: {1} Gain: {2}".format
				(
					self.TargetedAttributeValue,
					self.TargetedAttribute,
					self.InformationGain
				)
			]

		keys = self.Children.keys()
		keys.sort()
		for index, key in enumerate(keys):
			childList = self.Children[key]._to_string()
			for i in range(len(childList)):
				if (childList[i].startswith("|--- ") or
					childList[i].startswith("|    ") or
					childList[i].startswith("     ")):
					if index == len(keys) - 1:
						childList[i] = "     " + childList[i]
					else:
						childList[i] = "|    " + childList[i]
				else:
					childList[i] = "|--- " + childList[i]

			result.extend(childList)
			
		return result

	def process_record(self, record):
		# Terminal Case
		if (self.TargetedAttribute, self.InformationGain) == (None, None):
			# Select a random record from the trained dataset
			sampleRecord = random.choice(self.Dataset)
			
			# Return that record's classification
			return sampleRecord[1]

		# Save original record in case of error
		originalRecord = record[:]

		# Select the attribute to match on
		attrIndex = self.AttributeNames.index(self.TargetedAttribute)

		# Save that attribute's value
		attrValue = record[0][attrIndex]

		# Strip out that attribute
		record = (record[0][:attrIndex] + record[0][attrIndex+1:], record[1])

		if attrValue in self.Children:
			# Recursively call ProcessRecord on the DecisionTree associated with
			# the selected attribute's value and stuff some error catching that
			# will help display program flow
			try:
				return self.Children[attrValue].process_record(record)
			except Exception:
				print("record",originalRecord)
				print("self.Children.keys()",self.Children.keys())
				print("attrValue",attrValue)
				print("attrIndex",attrIndex)
				print("self.AttributeNames",self.AttributeNames)
				print("self.TargetedAttribute",self.TargetedAttribute)
				print("")
				print("")
				raise
		else:
			# Select a random record
			sampleRecord = random.choice(self.Dataset)
			return sampleRecord[1]

	def Size(self):
		sum = 1
		if len(self.Children) == 0:
			return 1
		for child in self.Children:
			sum += self.Children[child].Size()
		return sum
	
	def Leaves(self):
		sum = 0
		if len(self.Children) == 0:
			return 1
		for child in self.Children:
			sum += self.Children[child].Leaves()
		return sum